"""
DPA 线程时间分析绘图脚本 (优化版)
用于可视化不同线程配置下的性能数据
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import os
import matplotlib as mpl

# 设置全局字体大小
mpl.rcParams['font.size'] = 12

# ===== 配置类 =====
class PlotConfig:
    """绘图配置类，集中管理所有可配置参数"""
    def __init__(self):
        # 输入文件配置
        self.input_csv = 'dpa_thread_analyze_v5.csv'
        self.input_dir = '../../results_raw/dpa_case_study'
        
        # 输出文件配置
        self.output_dir = '../../plots/dpa_case_study'
        self.output_prefix = 'dpa_thread_analysis'
        
        # 绘图参数
        self.batch_size = 16777216  # 分析的批处理大小
        self.dpa_types = [0, 2]  # 分析的DPA类型
        self.thread_nums = [2, 4, 8, 16, 32, 64, 96, 128, 160, 192, 224, 253]  # 分析的线程数
        
        # 颜色配置
        self.colors = {
            'normal_mem_access': '#1f77b4',    # 蓝色
            'get_host_mem_addr': '#ff7f0e',    # 橙色
            'operation_time': '#2ca02c',       # 绿色
            'atomic_store_time': '#d62728',     # 红色
            'overhead': '#ffffff',             # 白色（损耗部分）
            'scheduling': '#9467bd',           # 紫色（调度时间）
            'execution': '#8c564b'             # 棕色（执行时间）
        }
        
        # 样式配置
        self.bar_width = 0.35
        self.figsize = (14, 10)  # 增加高度以提供更多空间
        self.dpi = 300
        self.fontsize = 12
        self.title_fontsize = 16
        
        # 高级配置
        self.show_plot = True
        self.save_plot = True
    
    def get_input_path(self):
        """获取输入文件路径"""
        return os.path.join(self.input_dir, self.input_csv)
    
    def get_output_path(self):
        """获取输出文件路径"""
        return os.path.join(self.output_dir, f"{self.output_prefix}_batch_{self.batch_size}.png")
    
    def print_config(self):
        """打印当前配置"""
        print("\n===== 绘图配置 =====")
        print(f"批处理大小: {self.batch_size}")
        print(f"DPA类型: {self.dpa_types}")
        print(f"线程数: {self.thread_nums}")
        print(f"输入文件: {self.get_input_path()}")
        print(f"输出文件: {self.get_output_path()}")
        print("====================\n")

# ===== 数据处理函数 =====
def load_and_filter_data(config):
    """加载并过滤数据"""
    # 加载数据
    input_path = config.get_input_path()
    data = pd.read_csv(input_path)
    
    # 筛选数据
    filtered_data = data[
        (data['batch_size'] == config.batch_size) &
        (data['dpa_type'].isin(config.dpa_types)) &
        (data['thread_num'].isin(config.thread_nums))
    ]
    
    # 按线程数和DPA类型排序
    filtered_data = filtered_data.sort_values(['thread_num', 'dpa_type'])
    
    return filtered_data

def calculate_metrics(data):
    """计算所需的指标"""
    data['p50_time'] = data['p50_time'] / 2
    # 计算线程时间占比（限制在0-100%）
    data['thread_ratio'] = (data['thread_time'] / data['p50_time']) * 100
    data['thread_ratio'] = data['thread_ratio'].clip(upper=100)
    
    # 计算执行时间（线程实际工作时间）
    data['execution_time'] = data['p50_time'] * data['thread_ratio'] / 100
    
    # 计算调度时间（总时间减去执行时间）
    data['scheduling_time'] = data['p50_time'] - data['execution_time']
    
    # 将四个组件时间转换为微秒
    components = ['normal_mem_access', 'get_host_mem_addr', 'operation_time', 'atomic_store_time']
    for comp in components:
        data[f'{comp}_us'] = data[comp] / 1800
    
    # 计算组件时间总和
    data['components_sum'] = data[[f'{comp}_us' for comp in components]].sum(axis=1)
    
    # 计算损耗时间（线程时间减去组件时间总和）
    data['overhead_us'] = data['thread_time'] - data['components_sum']
    data['overhead_us'] = data['overhead_us'].clip(lower=0)  # 确保非负
    
    return data

# ===== 绘图函数 =====
def create_legend(config):
    """创建自定义图例"""
    legend_patches = [
        mpatches.Patch(facecolor=config.colors['normal_mem_access'], edgecolor='black', label='Normal Mem Access'),
        mpatches.Patch(facecolor=config.colors['get_host_mem_addr'], edgecolor='black', label='Get Host Mem Addr'),
        mpatches.Patch(facecolor=config.colors['operation_time'], edgecolor='black', label='Operation Time'),
        mpatches.Patch(facecolor=config.colors['atomic_store_time'], edgecolor='black', label='Atomic Store Time'),
        mpatches.Patch(facecolor=config.colors['overhead'], edgecolor='black', label='Overhead'),
        mpatches.Patch(facecolor=config.colors['scheduling'], edgecolor='black', label='Scheduling Time'),
    ]
    return legend_patches

def plot_dpa_analysis(data, config):
    """创建DPA分析图"""
    # 设置图形（增加高度提供更多空间）
    fig, ax = plt.subplots(figsize=config.figsize, dpi=config.dpi)
    
    # 获取唯一线程数
    thread_nums = sorted(data['thread_num'].unique())
    dpa_types = sorted(data['dpa_type'].unique())
    
    # 计算柱状图位置
    x_positions = np.arange(len(thread_nums))
    
    # 定义堆叠顺序
    components = [
        'normal_mem_access_us', 
        'get_host_mem_addr_us', 
        'operation_time_us', 
        'atomic_store_time_us',
        'overhead_us'
    ]
    
    # 为每种DPA类型创建数据容器
    dpa_data = {}
    for dpa_type in dpa_types:
        dpa_data[dpa_type] = data[data['dpa_type'] == dpa_type].sort_values('thread_num')
    
    # 绘制堆叠柱状图
    for i, dpa_type in enumerate(dpa_types):
        # 计算位置偏移
        offset = i * config.bar_width
        
        # 重置底部位置
        bottom = np.zeros(len(thread_nums))
        
        # 绘制执行时间组件
        for comp in components:
            comp_times = dpa_data[dpa_type][comp].values
            
            # 设置颜色
            color = config.colors['overhead'] if comp == 'overhead_us' else config.colors[comp.split('_us')[0]]
            
            # 绘制柱状图
            ax.bar(
                x_positions + offset, comp_times, 
                width=config.bar_width, bottom=bottom,
                color=color, edgecolor='black', label=comp if i == 0 else None
            )
            
            # 更新底部位置
            bottom += comp_times
        
        # 绘制调度时间
        scheduling_times = dpa_data[dpa_type]['scheduling_time'].values
        ax.bar(
            x_positions + offset, scheduling_times, 
            width=config.bar_width, bottom=bottom,
            color=config.colors['scheduling'], edgecolor='black',
            label='Scheduling Time' if i == 0 else None
        )
        
        # 添加总时间标签
        for j, pos in enumerate(x_positions):
            total_time = dpa_data[dpa_type].iloc[j]['p50_time']
            ratio = dpa_data[dpa_type].iloc[j]['thread_ratio']
            ax.text(
                pos + offset, bottom[j] + scheduling_times[j] + 0.05 * max(scheduling_times),
                f'{total_time:.1f} us\n({ratio:.0f}%)', 
                ha='center', va='bottom', fontsize=config.fontsize - 2
            )
    
    # 设置图形属性
    ax.set_title(
        f'DPA Thread Time Analysis (Batch Size = {config.batch_size})', 
        fontsize=config.title_fontsize
    )
    ax.set_xlabel('Thread Number', fontsize=config.fontsize)
    ax.set_ylabel('Time (us)', fontsize=config.fontsize)
    group_centers = x_positions + config.bar_width * (len(dpa_types) - 1) / 2
    ax.set_xticks(group_centers)

    # 设置标签为正常字体，不旋转，居中对齐
    ax.set_xticklabels(thread_nums, fontsize=config.fontsize, rotation=0, ha='center')
    
    ax.grid(axis='y', linestyle='--', alpha=0.7)
    
    # 添加图例
    legend_patches = create_legend(config)
    ax.legend(
        handles=legend_patches, 
        loc='upper left', 
        fontsize=config.fontsize,
        ncol=2
    )
    
    # 添加DPA类型标签（使用注释线避免重叠）
    for i, dpa_type in enumerate(dpa_types):
        offset = i * config.bar_width - 0.5 * config.bar_width
        ax.annotate(
            f'DPA {dpa_type}',
            xy=(x_positions[-1] + offset + config.bar_width/2, 0),
            xytext=(0, -40 * (dpa_type + 1) ),  # 文本位置在x轴下方40点
            textcoords='offset points',
            ha='center', va='top', fontsize=config.fontsize,
            arrowprops=dict(arrowstyle="-", color='gray', alpha=0.5)
        )
    
    # 手动调整布局（替代tight_layout）
    plt.subplots_adjust(
        left=0.10,   # 左侧空间
        right=0.95,  # 右侧空间
        bottom=0.25, # 底部空间（增加更多空间）
        top=0.90     # 顶部空间
    )
    
    # 保存图像（不再使用bbox_inches='tight'）
    if config.save_plot:
        output_path = config.get_output_path()
        os.makedirs(config.output_dir, exist_ok=True)
        plt.savefig(output_path, dpi=config.dpi)
        print(f"图像已保存到: {output_path}")
    
    # 显示图像
    if config.show_plot:
        plt.show()

# ===== 主函数 =====
def main():
    """主函数"""
    # 初始化配置
    config = PlotConfig()
    config.print_config()
    
    # 加载并处理数据
    data = load_and_filter_data(config)
    data = calculate_metrics(data)
    
    # 绘制图像
    plot_dpa_analysis(data, config)

if __name__ == "__main__":
    main()
